package com.xiashitech.agent.instrumentation.redis.redisson.v317.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.instrumentation.bridge.Java8Bridge;
import com.xiashitech.agent.instrumentation.context.BaggageIsModify;
import com.xiashitech.agent.instrumentation.context.InstContext;
import com.xiashitech.agent.instrumentation.redis.RedisData;
import com.xiashitech.agent.instrumentation.redis.redisson.XSCompletableFuture;
import com.xiashitech.agent.instrumentation.redis.redisson.XSEndOperationListener;
import com.xiashitech.agent.instrumentation.redis.redisson.XSPromiseWrapper;
import com.xiashitech.agent.instrumentation.redis.redisson.XSRedissonRequest;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.utils.DurationPrinter;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import static com.xiashitech.agent.instrumentation.redis.redisson.v317.XSRedissonSingletons.instrumenter;
import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.named;

public class XSRedisConnectionInstrumentation implements TypeInstrumentation {
    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return named("org.redisson.client.RedisConnection");
    }

    @Override
    public void transform(TypeTransformer transformer) {
        transformer.applyAdviceToMethod(
                isMethod().and(named("send")), this.getClass().getName() + "$SendAdvice");
    }

    @SuppressWarnings("unused")
    public static class SendAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.Argument(0) Object argument,
                @Advice.Local("otelRequest") RedisData redisData,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("otelPromise") XSCompletableFuture xsCompletableFuture) {
            DurationPrinter durationPrinter = new DurationPrinter("XSRedisConnectionInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            XSRedissonRequest redissonRequest = new XSRedissonRequest();
            redissonRequest.setCommand(argument);
            XSPromiseWrapper promise = redissonRequest.getPromiseWrapper();
            if (promise == null) {
                return;
            }
            try {
                xsCompletableFuture = (XSCompletableFuture)promise;
                if(!redissonRequest.isSupportedCommand()) {
                    return;
                }
                redissonRequest.setBaggage(Java8Bridge.currentForBaggage());
                redisData = RedisData.builder().redissonRequest(redissonRequest).build();
                redisData.setRedisOpetion(redissonRequest.getOperation());
                redisData.setParams(redissonRequest.getALLParams());
                redissonRequest.setRedisDataParams(redisData);
                Context parentContext = currentContext();
                if (!instrumenter().shouldStart(parentContext, redisData)) {
                    return;
                }
                parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
                context = instrumenter().start(parentContext, redisData);
                context = BaggageUtil.syncBaggage().build().storeInContext(context);
                scope = context.makeCurrent();
                promise.setEndOperationListener(new XSEndOperationListener(instrumenter(), context, redisData));
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("SendAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.Thrown Throwable throwable,
                @Advice.Local("otelRequest") RedisData redisData,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("otelPromise") XSCompletableFuture xsCompletableFuture) {
            DurationPrinter durationPrinter = new DurationPrinter("XSRedisConnectionInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if (scope == null) {
                return;
            }
            scope.close();
            if (throwable != null) {
                instrumenter().end(context, redisData, null, throwable);
            }
            InstContext.set(xsCompletableFuture);
            BaggageIsModify.setIsModify(Boolean.FALSE);
            durationPrinter.printElapseDuration("SendAdvice e");
        }
    }
}
